#ifndef LA_EPILOGUE_MULTIPLY_SCALAR_ADD
#define LA_EPILOGUE_MULTIPLY_SCALAR_ADD

template<typename OutputDataType_, typename Enable_=void>
class BlockEpilogueMultiplyScalarAdd {
    static_assert(AscendC::Std::always_false_v<OutputDataType_>, "BlockEpilogueMultiplyScalarAdd only supports fp32 as OutputDataType");
};

template<typename OutputDataType_>
class BlockEpilogueMultiplyScalarAdd<
    OutputDataType_,
    AscendC::Std::enable_if_t<
        AscendC::Std::is_same_v<OutputDataType_, float>
    >
> {
public:
    using OutputDataType = OutputDataType_;

    struct Arguments {};
    struct Params {};

    __host_aicore__ static Params InitParams(Arguments const& args) { return { }; }

    __host_aicore__ static size_t GetWorkSpaceSize(Arguments const& args) { return 0; }

    __aicore__ inline void Init(Params const& params) { }
private:

public:
    __aicore__ inline void RunMultiplyScalar(
        const AscendC::LocalTensor<float>& outTemp,
        const AscendC::LocalTensor<float>& input1,
        const float param1,
        int32_t dataCount
    ) {
        if (dataCount == 0) {
            return ;
        }

        AscendC::PipeBarrier<PIPE_V>();
        AscendC::Muls(outTemp, input1, param1, dataCount);
        AscendC::PipeBarrier<PIPE_V>();
    }

    __aicore__ inline void RunAdd(
        const AscendC::LocalTensor<OutputDataType>& out,
        const AscendC::LocalTensor<float>& outTemp,
        const AscendC::LocalTensor<float>& param2,
        int32_t dataCount
    ) {
        if (dataCount == 0) {
            return ;
        }

        if constexpr (AscendC::Std::is_same_v<OutputDataType, float>) {
            AscendC::PipeBarrier<PIPE_V>();
            AscendC::Add(out, outTemp, param2, dataCount);
            AscendC::PipeBarrier<PIPE_V>();
        } else {
            AscendC::PipeBarrier<PIPE_V>();
            AscendC::Add(outTemp, outTemp, param2, dataCount);
            AscendC::PipeBarrier<PIPE_V>();
            AscendC::Cast(out, outTemp, AscendC::RoundMode::CAST_ROUND, dataCount);
            AscendC::PipeBarrier<PIPE_V>();
        }
    }
};

#endif // !LA_EPILOGUE_MULTIPLY_SCALAR_ADD